import { NextRequest, NextResponse } from 'next/server'
import { getServerSession } from 'next-auth'
import { authOptions } from '@/lib/auth'
import prisma from '@/lib/prisma'
import { z } from 'zod'
import {
  executeCodeInPiston,
  mapPistonResultToSimplified,
  PistonExecuteRequest,
  SimplifiedPistonStatus,
  MappedPistonResult,
} from '@/lib/pistonService'
import { TestCase, ExperimentStage } from '@prisma/client'

// Zod schema for answer payload
// 允许答案是各种可能的结构，或者为null
const answerSchema = z.record(z.any()).nullable();

const saveAnswerPayloadSchema = z.object({
  answer: answerSchema,
  evaluate: z.boolean().optional(),
});

interface StageContent {
  language?: string;
  languageVersion?: string;
  testCases?: Array<TestCase & { input?: string, expectedOutput?: string, points: number }>;
  timeLimitMs?: number;
  memoryLimitKb?: number;
  [key: string]: any; // Allow other properties
}

export async function POST(
  req: NextRequest,
  { params }: { params: { id: string; stageId: string } }
) {
  const session = await getServerSession(authOptions);

  if (!session?.user?.id || session.user.role !== 'STUDENT') {
    return NextResponse.json({ error: 'Forbidden. Only students can save answers.' }, { status: 403 });
  }

  const { id: experimentId, stageId } = params;
  const studentId = session.user.id;

  if (!experimentId || !stageId) {
    return NextResponse.json({ error: 'Experiment ID and Stage ID are required' }, { status: 400 });
  }

  let payload;
  try {
    const body = await req.json();
    payload = saveAnswerPayloadSchema.parse(body);
  } catch (error) {
    if (error instanceof z.ZodError) {
      return NextResponse.json({ error: 'Invalid input data for answer', details: error.errors }, { status: 400 });
    }
    return NextResponse.json({ error: 'Invalid JSON payload for answer' }, { status: 400 });
  }

  try {
    const stage = await prisma.experimentStage.findUnique({
      where: { id: stageId },
      include: { experiment: true }
    });

    if (!stage || stage.experimentId !== experimentId) {
      return NextResponse.json({ error: 'Stage not found or does not belong to this experiment.' }, { status: 404 });
    }
    
    const now = new Date();
    if (stage.experiment.startTime && now < new Date(stage.experiment.startTime)) {
        return NextResponse.json({ error: 'Experiment has not started yet.' }, { status: 403 });
    }
    if (stage.experiment.endTime && now > new Date(stage.experiment.endTime) && !stage.experiment.allowLateSubmission) {
        return NextResponse.json({ error: 'Experiment has ended and late submissions are not allowed.' }, { status: 403 });
    }
    
    let experimentSubmission = await prisma.experimentSubmission.findFirst({
      where: { studentId, experimentId },
    });

    if (!experimentSubmission) {
      experimentSubmission = await prisma.experimentSubmission.create({
        data: { studentId, experimentId, status: 'IN_PROGRESS', draftStageAnswers: JSON.stringify({}) },
      });
    }

    let currentDraftAnswersObject: Record<string, any>;
    try {
        currentDraftAnswersObject = experimentSubmission.draftStageAnswers 
                                    ? JSON.parse(experimentSubmission.draftStageAnswers) 
                                    : {};
    } catch (parseError) {
        console.error("Failed to parse draftStageAnswers JSON, resetting to empty object:", parseError);
        currentDraftAnswersObject = {}; // Reset if parsing fails
    }
    
    currentDraftAnswersObject[stageId] = { 
        answer: payload.answer || {}, 
        lastSaved: new Date().toISOString() 
    };

    const updatedDraft = await prisma.experimentSubmission.update({
      where: { id: experimentSubmission.id },
      data: { draftStageAnswers: JSON.stringify(currentDraftAnswersObject), updatedAt: new Date() },
    });

    if (stage.type === 'CODE' && payload.evaluate && payload.answer && typeof payload.answer === 'object' && 'code' in payload.answer) {
      console.log(`请求对阶段 ${stageId} 进行 Piston 即时评测。`);
      const studentCodeAnswer = (payload.answer as { code: string }).code;
      const studentSelectedLanguageInfo = (payload.answer as { language?: string, version?: string });
      
      let evaluationResults: any[] = [];
      let overallFeedback = '评测已执行。';

      try {
        const stageContent: StageContent = typeof stage.content === 'string' ? JSON.parse(stage.content) : stage.content;
        
        const pistonLanguage = studentSelectedLanguageInfo?.language || stageContent?.language;
        const pistonVersion = studentSelectedLanguageInfo?.version || stageContent?.languageVersion;

        if (!studentCodeAnswer) throw new Error('未提供评测代码。');
        if (!pistonLanguage) throw new Error('无法确定评测语言。');
        if (!pistonVersion) throw new Error(`语言 ${pistonLanguage} 未指定版本。Piston 需要版本信息。`);

        const fileName = `main.${pistonLanguage.split(/[^a-zA-Z0-9]/)[0].toLowerCase() || 'txt'}`;

        const testCasesFromDB = stageContent?.testCases || [];
        
        if (testCasesFromDB.length === 0) {
          overallFeedback = '阶段未配置测试用例，无法执行即时评测。';
        } else {
          for (const tc of testCasesFromDB) {
            if (!tc.id) { 
              console.warn(`Piston 即时评测：跳过测试用例，缺少DB ID: ${tc.name}`); 
              evaluationResults.push({ 
                testCaseId: tc.name,
                name: tc.name, 
                status: 'INTERNAL_ERROR', 
                runtimeError: '测试用例配置错误 (缺少ID)',
                isCorrect: false,
                pointsAwarded: 0,
              }); 
              continue; 
            }
            
            const pistonPayload: PistonExecuteRequest = {
              language: pistonLanguage,
              version: pistonVersion,
              files: [{ name: fileName, content: studentCodeAnswer }],
              stdin: tc.input || '',
              args: [],
              compile_timeout: stageContent.compile_timeout || 10000,
              run_timeout: stageContent.run_timeout || stageContent.timeLimitMs || 3000,
              compile_memory_limit: stageContent.compile_memory_limit || -1,
              run_memory_limit: stageContent.memoryLimitKb ? stageContent.memoryLimitKb * 1024 : -1,
            };

            let mappedResult: MappedPistonResult;

            try {
              console.log(`Piston 即时评测：为阶段 ${stageId} 测试用例 ${tc.id} (${tc.name}) 提交到 Piston`);
              const pistonResponse = await executeCodeInPiston(pistonPayload);
              mappedResult = mapPistonResultToSimplified(pistonResponse, tc.expectedOutput, tc.points);

            } catch (pistonError: any) {
              console.error(`Piston 即时评测失败 for TC ${tc.name}:`, pistonError);
              let errorMsg = pistonError.message || 'Piston 服务通信失败';
              if (pistonError.message && pistonError.message.toLowerCase().includes('rate limit')) {
                errorMsg = 'Piston 评测请求过于频繁，请稍后再试。';
              }
               if (pistonError.message && (pistonError.message.toLowerCase().includes('unknown language') || pistonError.message.toLowerCase().includes('unknown version'))){
                errorMsg = `Piston 错误: 未知的语言 (${pistonLanguage}) 或版本 (${pistonVersion}). 请检查实验配置。`;
              }
              mappedResult = {
                status: 'INTERNAL_ERROR',
                isCorrect: false,
                pointsAwarded: 0,
                actualOutput: null,
                runtimeError: errorMsg,
              };
            }
            evaluationResults.push({
              testCaseId: tc.id,
              name: tc.name,
              status: mappedResult.status,
              isCorrect: mappedResult.isCorrect,
              pointsAwarded: mappedResult.pointsAwarded,
              expectedOutput: tc.expectedOutput,
              actualOutput: mappedResult.actualOutput,
              runtimeError: mappedResult.runtimeError,
            });

            if (testCasesFromDB.indexOf(tc) < testCasesFromDB.length - 1) { 
                 await new Promise(resolve => setTimeout(resolve, 200));
             }
          }
        }
      } catch (evalError: any) {
        console.error('Piston 即时评测过程中发生错误:', evalError);
        overallFeedback = `评测失败: ${evalError.message}`;
        
        if (evalError.message && (evalError.message.toLowerCase().includes('无法确定评测语言') || evalError.message.toLowerCase().includes('未指定版本'))) {
            evaluationResults.push({
                testCaseId: 'config_error',
                name: '配置错误',
                status: 'INTERNAL_ERROR',
                runtimeError: evalError.message,
                isCorrect: false,
                pointsAwarded: 0,
            });
        }
      }
      
      return NextResponse.json({ 
        message: '草稿已保存。' + overallFeedback,
        draftSubmission: { 
            experimentSubmissionId: updatedDraft.id,
                stageId: stageId,
            answerSnapshot: currentDraftAnswersObject[stageId] 
            },
        evaluationResults: evaluationResults
        });
    }

    return NextResponse.json({ 
        message: 'Draft answer saved successfully', 
        draftSubmission: { 
            experimentSubmissionId: updatedDraft.id,
            stageId: stageId,
            answerSnapshot: currentDraftAnswersObject[stageId] 
        }
    });

  } catch (error: any) {
    console.error('Error saving stage draft answer:', error);
    return NextResponse.json({ error: 'Internal server error while saving draft answer' , details: error.message}, { status: 500 });
  }
}